package other

import (
	"fmt"
	"gitee.com/jn-qq/similarity"
	"maps"
	"vector/tools"
)

type RAGSimilarity struct {
	baseResult string
	ragResult  string
	pollResult string
	resultData map[string]any
}

func (rs *RAGSimilarity) Reset() {
	rs.baseResult = ""
	rs.ragResult = ""
	rs.pollResult = ""
	clear(rs.resultData)
}

func (rs *RAGSimilarity) NewTaskData(row map[string]any) error {
	// 复制参数，避免修改原参数
	rs.resultData = maps.Clone(row)

	fields := map[string]*string{
		"参考答案":  &rs.baseResult,
		"RAG结果": &rs.ragResult,
		"润色结果":  &rs.pollResult,
	}

	for key, field := range fields {
		if txt, ok := row[key]; ok {
			*field = txt.(string)
		} else {
			rs.resultData["备注"] = "参数错误"
			return fmt.Errorf("参数错误")
		}
	}

	return nil
}

func (rs *RAGSimilarity) Do() {
	if (rs.ragResult == "" || rs.ragResult == "/") && (rs.pollResult == "" || rs.pollResult == "/") {
		rs.resultData["RAG相似度"] = 0
		rs.resultData["润色相似度"] = 0
		return
	}
	source, err := tools.SIM.TextSimilar(similarity.Cosine, rs.baseResult, rs.ragResult, rs.pollResult)
	if err != nil {
		rs.resultData["备注"] = err.Error()
		return
	}
	rs.resultData["RAG相似度"] = source[0]
	rs.resultData["润色相似度"] = source[1]
}

func (rs *RAGSimilarity) Result() map[string]any {
	// 复制结果，避免对象引用被修改
	return maps.Clone(rs.resultData)
}
